#!/usr/bin/python
# -*-coding:utf-8-*-

'''
定义一些时间序列算子
'''

import os
import numpy as np
import pandas as pd
from joblib import wrap_non_picklable_objects

from gplearn_evolve.base.utils import d_type_exam
from gplearn_evolve.base.utils import d_map

from gplearn_evolve.base.utils import keys_core_pd
from gplearn_evolve.base.utils import keys_core_pd_v2
from gplearn_evolve.base.utils import np_neu_keys_to_pd

from statsmodels.regression.rolling import RollingOLS
import statsmodels.api as sm

# 1
@wrap_non_picklable_objects
def def_keyfunc_cs_zscore(keys, key_by_path=True):
    def gp_cs_zscore(X):
        if key_by_path:
            # load key
            load_keys = np.load(keys,
                                mmap_mode='r')
        else:
            load_keys = keys.copy()

        if len(load_keys) != len(X):
            return X

        if not key_by_path:
            values = keys_core_pd(keys, X)
        else:
            values = keys_core_pd_v2(load_keys, X)

        values = values.groupby(as_index=False,
                                level='date',
                                sort=False).apply(lambda x:(x - x.mean())/x.std())#return NaN if std is zero

        # TODO - 将inf设置为nan
        values[np.isinf(values['f1'])] = np.nan

        del load_keys

        return values['f1'].values

    return gp_cs_zscore

# 2
@wrap_non_picklable_objects
def def_keyfunc_cs_rank(keys, key_by_path=True):
    def gp_cs_rank(X):
        if key_by_path:
            # load key
            load_keys = np.load(keys,
                                mmap_mode='r')
        else:
            load_keys = keys.copy()

        if len(load_keys) != len(X):
            return X

        if not key_by_path:
            values = keys_core_pd(keys, X)
        else:
            values = keys_core_pd_v2(load_keys, X)

        values = values.groupby(as_index=False,
                                level='date',
                                sort=False).rank(pct=True)

        # TODO - 将inf设置为nan
        values[np.isinf(values['f1'])] = np.nan

        del load_keys

        return values['f1'].values

    return gp_cs_rank

# 3
@wrap_non_picklable_objects
def def_keyfunc_cs_minmax(keys, key_by_path=True):
    def gp_minmax(X):
        if key_by_path:
            # load key
            load_keys = np.load(keys,
                                mmap_mode='r')
        else:
            load_keys = keys.copy()

        if len(load_keys) != len(X):
            return X

        if not key_by_path:
            values = keys_core_pd(keys, X)
        else:
            values = keys_core_pd_v2(load_keys, X)

        values = values.groupby(as_index=False,
                                level='date',
                                sort=False).apply(lambda x:(x-x.min())/(x.max()-x.min()))

        # TODO - 将inf设置为nan
        values[np.isinf(values['f1'])] = np.nan

        del load_keys

        return values['f1'].values

    return gp_minmax

# 4
@wrap_non_picklable_objects
def def_keyfunc_ts_delay(keys, base=5, key_by_path=True):
    def gp_ts_delay(X, d):
        if key_by_path:
            # load key
            load_keys = np.load(keys,
                                mmap_mode='r')
        else:
            load_keys = keys.copy()

        if len(load_keys) != len(X):
            return X

        params = d_map(d_type_exam(d), base=base)

        if not key_by_path:
            values = keys_core_pd(keys, X)
        else:
            values = keys_core_pd_v2(load_keys, X)

        values = values.groupby(as_index=False,
                                level='stock_code',
                                sort=False).shift(params)

        # TODO - 将inf设置为nan
        values[np.isinf(values['f1'])] = np.nan

        del load_keys

        return values['f1'].values

    return gp_ts_delay

# 5
@wrap_non_picklable_objects
def def_keyfunc_ts_corr(keys, base=5, key_by_path=True):
    def gp_ts_corr(X, Y, d):
        if key_by_path:
            # load key
            load_keys = np.load(keys,
                                mmap_mode='r')
        else:
            load_keys = keys.copy()

        if len(load_keys) != len(X):
            return X

        params = d_map(d_type_exam(d), base=base)

        if not key_by_path:
            values = keys_core_pd(keys, X, Y)
        else:
            values = keys_core_pd_v2(load_keys, X, Y)

        values = values.groupby(as_index=False,
                                level='stock_code',
                                sort=False).apply(lambda x:x.iloc[:,0].rolling(window=params,
                                                                                 min_periods=params//2).corr(x.iloc[:,1]))

        # TODO - 将inf设置为nan
        # values[np.isinf(values['f1'])] = np.nan
        values[np.isinf(values)] = np.nan

        del load_keys

        return values.values

    return gp_ts_corr

# 6
@wrap_non_picklable_objects
def def_keyfunc_ts_min(keys, base=5, key_by_path=True):
    def gp_ts_min(X, d):
        if key_by_path:
            # load key
            load_keys = np.load(keys,
                                mmap_mode='r')
        else:
            load_keys = keys.copy()

        if len(load_keys) != len(X):
            return X

        params = d_map(d_type_exam(d), base=base)

        if not key_by_path:
            values = keys_core_pd(keys, X)
        else:
            values = keys_core_pd_v2(load_keys, X)

        values = values.groupby(as_index=False,
                                level='stock_code',
                                sort=False).apply(lambda x:x.rolling(window=params, min_periods=params//2).min())

        # TODO - 将inf设置为nan
        values[np.isinf(values['f1'])] = np.nan

        del load_keys

        return values['f1'].values

    return gp_ts_min

# 7
@wrap_non_picklable_objects
def def_keyfunc_ts_max(keys, base=5, key_by_path=True):
    def gp_ts_max(X, d):
        if key_by_path:
            # load key
            load_keys = np.load(keys,
                                mmap_mode='r')
        else:
            load_keys = keys.copy()

        if len(load_keys) != len(X):
            return X

        params = d_map(d_type_exam(d), base=base)

        if not key_by_path:
            values = keys_core_pd(keys, X)
        else:
            values = keys_core_pd_v2(load_keys, X)

        values = values.groupby(as_index=False,
                                level='stock_code',
                                sort=False).apply(lambda x:x.rolling(window=params,
                                                                     min_periods=params//2).max())

        # TODO - 将inf设置为nan
        values[np.isinf(values['f1'])] = np.nan

        del load_keys

        return values['f1'].values

    return gp_ts_max

# 8
@wrap_non_picklable_objects
def def_keyfunc_ts_kurt(keys, base=5, key_by_path=True):
    def gp_ts_kurt(X, d):
        if key_by_path:
            # load key
            load_keys = np.load(keys,
                                mmap_mode='r')
        else:
            load_keys = keys.copy()

        if len(load_keys) != len(X):
            return X

        params = d_map(d_type_exam(d), base=base)

        if not key_by_path:
            values = keys_core_pd(keys, X)
        else:
            values = keys_core_pd_v2(load_keys, X)

        values = values.groupby(as_index=False,
                                level='stock_code',
                                sort=False).apply(lambda x:x.rolling(window=params,
                                                                     min_periods=params//2).kurt())

        # TODO - 将inf设置为nan
        values[np.isinf(values['f1'])] = np.nan

        del load_keys

        return values['f1'].values

    return gp_ts_kurt

# 9
@wrap_non_picklable_objects
def def_keyfunc_ts_skew(keys, base=5, key_by_path=True):
    def gp_ts_skew(X, d):
        if key_by_path:
            # load key
            load_keys = np.load(keys,
                                mmap_mode='r')
        else:
            load_keys = keys.copy()

        if len(load_keys) != len(X):
            return X

        params = d_map(d_type_exam(d), base=base)

        if not key_by_path:
            values = keys_core_pd(keys, X)
        else:
            values = keys_core_pd_v2(load_keys, X)

        values = values.groupby(as_index=False,
                                level='stock_code',
                                sort=False).apply(lambda x:x.rolling(window=params,
                                                                     min_periods=params//2).skew())

        # TODO - 将inf设置为nan
        values[np.isinf(values['f1'])] = np.nan

        del load_keys

        return values['f1'].values

    return gp_ts_skew

# 10
@wrap_non_picklable_objects
def def_keyfunc_ts_std(keys, base=5, key_by_path=True):
    def gp_ts_std(X, d):
        if key_by_path:
            # load key
            load_keys = np.load(keys,
                                mmap_mode='r')
        else:
            load_keys = keys.copy()

        if len(load_keys) != len(X):
            return X

        params = d_map(d_type_exam(d), base=base)

        if not key_by_path:
            values = keys_core_pd(keys, X)
        else:
            values = keys_core_pd_v2(load_keys, X)

        values = values.groupby(as_index=False,
                                level='stock_code',
                                sort=False).apply(lambda x:x.rolling(window=params,
                                                                     min_periods=params//2).std())

        # TODO - 将inf设置为nan
        values[np.isinf(values['f1'])] = np.nan

        del load_keys

        return values['f1'].values

    return gp_ts_std

# 11
@wrap_non_picklable_objects
def def_keyfunc_ts_sum(keys, base=5, key_by_path=True):
    def gp_ts_sum(X, d):
        if key_by_path:
            # load key
            load_keys = np.load(keys,
                                mmap_mode='r')
        else:
            load_keys = keys.copy()

        if len(load_keys) != len(X):
            return X

        params = d_map(d_type_exam(d), base=base)

        if not key_by_path:
            values = keys_core_pd(keys, X)
        else:
            values = keys_core_pd_v2(load_keys, X)

        values = values.groupby(as_index=False,
                                level='stock_code',
                                sort=False).apply(lambda x:x.rolling(window=params,
                                                                       min_periods=params//2).sum())

        # TODO - 将inf设置为nan
        values[np.isinf(values['f1'])] = np.nan

        del load_keys

        return values['f1'].values

    return gp_ts_sum

# 12
@wrap_non_picklable_objects
def def_keyfunc_ts_norm(keys, base=5, key_by_path=True):
    def gp_ts_norm(X, d):
        if key_by_path:
            # load key
            load_keys = np.load(keys,
                                mmap_mode='r')
        else:
            load_keys = keys.copy()

        if len(load_keys) != len(X):
            return X

        params = d_map(d_type_exam(d), base=base)

        if not key_by_path:
            values = keys_core_pd(keys, X)
        else:
            values = keys_core_pd_v2(load_keys, X)

        values_mean = values.groupby(as_index=False,
                                level='stock_code',
                                sort=False).apply(lambda x:x.rolling(window=params,
                                                                     min_periods=params//2).mean())

        values_std = values.groupby(as_index=False,
                                level='stock_code',
                                sort=False).apply(lambda x:x.rolling(window=params,
                                                                     min_periods=params//2).std())

        values = (values - values_mean) / values_std

        # TODO - 将inf设置为nan
        values[np.isinf(values['f1'])] = np.nan

        del load_keys

        return values['f1'].values

    return gp_ts_norm

# 13
@wrap_non_picklable_objects
def def_keyfunc_ts_minmax(keys, base=5, key_by_path=True):
    def gp_ts_minmax(X, d):
        if key_by_path:
            # load key
            load_keys = np.load(keys,
                                mmap_mode='r')
        else:
            load_keys = keys.copy()

        if len(load_keys) != len(X):
            return X

        params = d_map(d_type_exam(d), base=base)

        if not key_by_path:
            values = keys_core_pd(keys, X)
        else:
            values = keys_core_pd_v2(load_keys, X)

        values_min = values.groupby(as_index=False,
                                level='stock_code',
                                sort=False).apply(lambda x:x.rolling(window=params,
                                                                       min_periods=params//2).min())

        values_max = values.groupby(as_index=False,
                                level='stock_code',
                                sort=False).apply(lambda x:x.rolling(window=params,
                                                                       min_periods=params//2).max())

        values = (values - values_min) / (values_max - values_min)

        # TODO - 将inf设置为nan
        values[np.isinf(values['f1'])] = np.nan

        del load_keys

        return values['f1'].values

    return gp_ts_minmax


'''中性化处理'''
# TODO - 市值中性标准化处理
# 14
@wrap_non_picklable_objects
def def_keyfunc_cs_cap_neu_rank(keys, neu_keys, cap_name='cap', group_num=10, key_by_path=True):
    def gp_cs_cap_neu_rank(X):
        if key_by_path:
            # load key
            load_keys = np.load(keys,
                                mmap_mode='r')

            load_neu_keys = np.load(neu_keys,
                                    mmap_mode='r')

            load_neu_keys = np_neu_keys_to_pd(load_neu_keys)
        else:
            load_keys = keys.copy()
            load_neu_keys = neu_keys.copy()

        if len(keys)!=len(X):
            return X

        if not key_by_path:
            values = keys_core_pd(keys, X)
        else:
            values = keys_core_pd_v2(load_keys, X)

        values = pd.concat([values, load_neu_keys], axis=1, join_axes=[values.index])

        group = values.groupby('date',
                               as_index=False,
                               sort=False)[cap_name].apply(pd.qcut, group_num, labels=False)

        group.index = group.index.droplevel(level=0)

        values['group'] = group.loc[values.index]

        neu_label = values.groupby(['date', 'group'],
                               as_index=False,
                               sort=False)[['f1']].rank(pct=True)

        neu_label = neu_label.loc[values.index]

        # TODO - 将inf设置为nan
        neu_label[np.isinf(neu_label['f1'])] = np.nan

        del load_keys, load_neu_keys

        return neu_label['f1'].values

    return gp_cs_cap_neu_rank

# 15
@wrap_non_picklable_objects
def def_keyfunc_cs_cap_neu_norm(keys, neu_keys, cap_name='cap', group_num=10, key_by_path=True):
    def gp_cs_cap_neu_norm(X):
        if key_by_path:
            # load key
            load_keys = np.load(keys,
                                mmap_mode='r')

            load_neu_keys = np.load(neu_keys,
                                    mmap_mode='r')

            load_neu_keys = np_neu_keys_to_pd(load_neu_keys)
        else:
            load_keys = keys.copy()
            load_neu_keys = neu_keys.copy()

        if len(keys)!=len(X):
            return X

        if not key_by_path:
            values = keys_core_pd(keys, X)
        else:
            values = keys_core_pd_v2(load_keys, X)

        values = pd.concat([values, load_neu_keys], axis=1, join_axes=[values.index])

        group = values.groupby('date',
                               as_index=False,
                               sort=False)[cap_name].apply(pd.qcut, group_num, labels=False)

        group.index = group.index.droplevel(level=0)

        values['group'] = group.loc[values.index]

        neu_label = values.groupby(['date', 'group'],
                               as_index=False,
                               sort=False)[['f1']].apply(lambda x: (x - x.mean())/ x.std())

        neu_label = neu_label.loc[values.index]

        # TODO - 将inf设置为nan
        neu_label[np.isinf(neu_label['f1'])] = np.nan

        del load_keys, load_neu_keys

        return neu_label['f1'].values

    return gp_cs_cap_neu_norm

# 16
@wrap_non_picklable_objects
def def_keyfunc_cs_cap_neu_demean(keys, neu_keys, cap_name='cap', group_num=10, key_by_path=True):
    def gp_cs_cap_neu_demean(X):
        if key_by_path:
            # load key
            load_keys = np.load(keys,
                                mmap_mode='r')

            load_neu_keys = np.load(neu_keys,
                                    mmap_mode='r')

            load_neu_keys = np_neu_keys_to_pd(load_neu_keys)
        else:
            load_keys = keys.copy()
            load_neu_keys = neu_keys.copy()

        if len(keys)!=len(X):
            return X

        if not key_by_path:
            values = keys_core_pd(keys, X)
        else:
            values = keys_core_pd_v2(load_keys, X)

        values = pd.concat([values, load_neu_keys], axis=1, join_axes=[values.index])

        group = values.groupby('date',
                               as_index=False,
                               sort=False)[cap_name].apply(pd.qcut, group_num, labels=False)

        group.index = group.index.droplevel(level=0)

        values['group'] = group.loc[values.index]

        neu_label = values.groupby(['date', 'group'],
                               as_index=False,
                               sort=False)[['f1']].apply(lambda x: x - x.mean())

        neu_label = neu_label.loc[values.index]

        # TODO - 将inf设置为nan
        neu_label[np.isinf(neu_label['f1'])] = np.nan

        del load_keys, load_neu_keys

        return neu_label['f1'].values

    return gp_cs_cap_neu_demean

# TODO - 行业中性标准化处理
# 17
@wrap_non_picklable_objects
def def_keyfunc_cs_ci1_neu_rank(keys, neu_keys, key_by_path=True):
    def gp_cs_ci1_neu_rank(X):
        if key_by_path:
            # load key
            load_keys = np.load(keys,
                                mmap_mode='r')

            load_neu_keys = np.load(neu_keys,
                                    mmap_mode='r')

            load_neu_keys = np_neu_keys_to_pd(load_neu_keys)
        else:
            load_keys = keys.copy()
            load_neu_keys = neu_keys.copy()

        if len(keys)!=len(X):
            return X

        if not key_by_path:
            values = keys_core_pd(keys, X)
        else:
            values = keys_core_pd_v2(load_keys, X)

        values = pd.concat([values, load_neu_keys], axis=1, join_axes=[values.index])

        neu_label = values.groupby(['date', 'ci1_code'],
                               as_index=False,
                               sort=False)[['f1']].rank(pct=True)

        neu_label = neu_label.loc[values.index]

        # TODO - 将inf设置为nan
        neu_label[np.isinf(neu_label['f1'])] = np.nan

        del load_keys, load_neu_keys

        return neu_label['f1'].values

    return gp_cs_ci1_neu_rank

# 18
@wrap_non_picklable_objects
def def_keyfunc_cs_ci1_neu_norm(keys, neu_keys, key_by_path=True):
    def gp_cs_ci1_neu_norm(X):
        if key_by_path:
            # load key
            load_keys = np.load(keys,
                                mmap_mode='r')

            load_neu_keys = np.load(neu_keys,
                                    mmap_mode='r')

            load_neu_keys = np_neu_keys_to_pd(load_neu_keys)
        else:
            load_keys = keys.copy()
            load_neu_keys = neu_keys.copy()

        if len(keys)!=len(X):
            return X

        if not key_by_path:
            values = keys_core_pd(keys, X)
        else:
            values = keys_core_pd_v2(load_keys, X)

        values = pd.concat([values, load_neu_keys], axis=1, join_axes=[values.index])

        neu_label = values.groupby(['date', 'ci1_code'],
                               as_index=False,
                               sort=False)[['f1']].apply(lambda x: (x - x.mean())/ x.std())

        neu_label = neu_label.loc[values.index]

        # TODO - 将inf设置为nan
        neu_label[np.isinf(neu_label['f1'])] = np.nan

        del load_keys, load_neu_keys

        return neu_label['f1'].values

    return gp_cs_ci1_neu_norm

# 19
@wrap_non_picklable_objects
def def_keyfunc_cs_ci1_neu_demean(keys, neu_keys, key_by_path=True):
    def gp_cs_ci1_neu_demean(X):
        if key_by_path:
            # load key
            load_keys = np.load(keys,
                                mmap_mode='r')

            load_neu_keys = np.load(neu_keys,
                                    mmap_mode='r')

            load_neu_keys = np_neu_keys_to_pd(load_neu_keys)
        else:
            load_keys = keys.copy()
            load_neu_keys = neu_keys.copy()

        if len(keys)!=len(X):
            return X

        if not key_by_path:
            values = keys_core_pd(keys, X)
        else:
            values = keys_core_pd_v2(load_keys, X)

        values = pd.concat([values, load_neu_keys], axis=1, join_axes=[values.index])

        neu_label = values.groupby(['date', 'ci1_code'],
                               as_index=False,
                               sort=False)[['f1']].apply(lambda x: x - x.mean())

        neu_label = neu_label.loc[values.index]

        # TODO - 将inf设置为nan
        neu_label[np.isinf(neu_label['f1'])] = np.nan

        del load_keys, load_neu_keys

        return neu_label['f1'].values

    return gp_cs_ci1_neu_demean


'''rolling ols'''
##
# @wrap_non_picklable_objects
# def def_keyfunc_ts_ols_residual(keys, base=5):
#     def gp_ts_ols_residual(X, Y, d):
#         if len(keys)!=len(X):
#             return X
#
#         params = d_map(d_type_exam(d), base=base)
#
#         values = keys_core_pd(keys, X)
#
#         def rolling_ols(df, params):
#             try:
#                 endog = df.iloc[:, 0]
#                 exog = sm.add_constant(df.iloc[:, 1])
#                 mod = RollingOLS(endog=endog,
#                                  exog=exog,
#                                  window=params,
#                                  min_nobs=params // 2,
#                                  missing='skip')
#
#                 model_res = mod.fit(params_only=True)
#
#                 df['f1'] = endog - (model_res.params * exog).sum(axis=1)
#
#                 return df['f1']
#             except:
#                 ret_df = df.copy()
#                 ret_df.iloc[:, :] = np.nan
#
#                 return ret_df['f1']
#
#         res = values.groupby(as_index=False,
#                              level='stock_code',
#                              sort=False).apply(rolling_ols, params)
#
#         # TODO - 将inf设置为nan
#         values[np.isinf(values['f1'])] = np.nan
#
#         return values['f1'].values
#
#     return gp_ts_ols_residual

# @wrap_non_picklable_objects
# def def_keyfunc_ts_ols_rsq(keys, base=5):
#     def gp_ts_ols_rsq(X, Y, d):
#         if len(keys)!=len(X):
#             return X
#
#         params = d_map(d_type_exam(d), base=base)
#
#         values = keys_core_pd(keys, X)
#
#         def rolling_ols(df, params):
#             try:
#                 endog = df.iloc[:, 0]
#                 exog = sm.add_constant(df.iloc[:, 1])
#                 mod = RollingOLS(endog=endog,
#                                  exog=exog,
#                                  window=params,
#                                  min_nobs=params // 2,
#                                  missing='skip')
#
#                 model_res = mod.fit(params_only=False)
#
#                 return model_res.rsquared_adj
#             except:
#                 ret_df = df.copy()
#                 ret_df.iloc[:, :] = np.nan
#
#                 return ret_df['f1']
#
#         rsquared_adj = values.groupby(as_index=False,
#                                       level='stock_code',
#                                       sort=False).apply(rolling_ols, params)
#
#         # TODO - 将inf设置为nan
#         values[np.isinf(values['f1'])] = np.nan
#
#         return values['f1'].values
#
#     return gp_ts_ols_rsq
